#include "arm_moveit_kinematics/arm_allfunction.h"

int main(int argc, char *argv[])
{       
        ros::init(argc, argv, "arm_allfunction_node");
        ros::NodeHandle allfunction_node;
        std::string group = "arm";
        std::string baseframe = "base_link";
        std::string yaml_path;
        // std::string yaml_path = "/home/clang/robotlab/robotarm/clang/21_tigerarm_ros/src/arm_moveit_kinematics/config/poses.yaml";
        arm_allfunction allfunc(group, baseframe);

        ros::param::get("~poses_yaml", yaml_path);
        allfunc.init(yaml_path);
        ros::AsyncSpinner spinner(3); 
        spinner.start();

        allfunc.kernel();
        
        ros::waitForShutdown();
        return 0;
}

arm_allfunction::arm_allfunction(std::string group_name, std::string baseframe_name):
groupName(group_name),
baseframeName(baseframe_name),
group(group_name), 
visual_tools(baseframe_name)
{
        joint_state_sub = allfunc_node.subscribe("/joint_states", 1000, &arm_allfunction::jointStateCallback, this);
        key_ctrl_sub = allfunc_node.subscribe("arm_keyctrl", 100, &arm_allfunction::keyCtrlCallback, this);
        joint_data_pub = allfunc_node.advertise<arm_moveit_kinematics::arm_data>("joint_serial_data", 100);
}

arm_allfunction::~arm_allfunction(){}

void arm_allfunction::keyCtrlCallback(const std_msgs::Int32 &msg)
{
        key_ctrl_msg.data = msg.data;
        ROS_INFO(":%d\r\n", key_ctrl_msg.data);
        is_press_key =1;
}

void arm_allfunction::jointStateCallback(const sensor_msgs::JointStateConstPtr  &msg)
{
        arm_moveit_kinematics::arm_data joint_data;
        joint_state_msg = *msg;
        // for(int i = 0; i<6; i++)
        // {
        //         ROS_INFO("%s: %.2f", joint_state_msg.name[i].c_str(), joint_state_msg.position[i]);
        // }
        joint_data.q0 =  joint_state_msg.position[0];
        joint_data.q1 =  joint_state_msg.position[1];
        joint_data.q2 =  joint_state_msg.position[2];
        joint_data.q3 =  joint_state_msg.position[3];
        joint_data.q4 =  joint_state_msg.position[4];
        joint_data.q5 =  joint_state_msg.position[5];

        joint_data_pub.publish(joint_data);
}

void arm_allfunction::setTarget(const geometry_msgs::Pose &msg)
{
        pose = msg;
        group.setPoseTarget(msg);
}

void arm_allfunction::kernel()
{
        while (ros::ok())
        {
                switch (key_ctrl_msg.data)
                {
                        case  KeyEnumdef::key_e:
                                setTarget(get_yaml_info.getYamlPose("pre_place"));
                                printPoseInfo();
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_r:
                                setTarget(get_yaml_info.getYamlPose("place"));
                                printPoseInfo();
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_t:
                               setTarget(get_yaml_info.getYamlPose("pick_island"));
                               memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                               printPoseInfo();
                                break;
                        case KeyEnumdef::key_q:
                                setTarget(get_yaml_info.getYamlPose("drive2"));
                                printPoseInfo();
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_w:
                                updatePose(ChgDirectionEnumdef::add, PoseDimensionEnumdef::z);
                                setTarget(pose);
                                printPoseInfo();
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_s:
                                updatePose(ChgDirectionEnumdef::minus, PoseDimensionEnumdef::z);
                                setTarget(pose);
                                printPoseInfo();
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_a:
                                updatePose(ChgDirectionEnumdef::add,PoseDimensionEnumdef::roll);
                                setTarget(pose);
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_d:
                                updatePose(ChgDirectionEnumdef::minus, PoseDimensionEnumdef::pitch);
                                setTarget(pose);         
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));                      
                                break;
                        case KeyEnumdef::key_W:
                                time_scale +=0.1;
                                ROS_INFO("time scale: %.2f", time_scale);
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_S:
                                time_scale -=0.1; 
                               ROS_INFO("time scale: %.2f", time_scale);
                               memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        case KeyEnumdef::key_l:
                                is_multiple_tp = 1;
                                multy_plan.trajectory_ =  combineTrajectory(get_yaml_info.getYamlPose("drive2"), get_yaml_info.getYamlPose("place"), &group);
                                memset(&key_ctrl_msg, 0 , sizeof(key_ctrl_msg));
                                break;
                        default: break;
                }
                
                if(is_press_key)
                {                 
                        is_press_key = !is_press_key;    
                        if(is_multiple_tp)
                        {
                                is_multiple_tp = !is_multiple_tp;
                                execute(multy_plan);                               
                        }
                        else
                        {
                                plan();
                                my_plan.trajectory_ = turnPlanTime(my_plan.trajectory_,  time_scale);
                                visualizePose(pose);
                                execute(my_plan);                  
                                moveit::core::RobotStatePtr start_state(group.getCurrentState());
                                group.setStartStateToCurrentState();             
                        }                      
                }
                 
        }
}

moveit::planning_interface::MoveItErrorCode arm_allfunction::plan()
{
        moveit::planning_interface::MoveItErrorCode success = group.plan(my_plan);
        return success;
}

void arm_allfunction::execute(moveit::planning_interface::MoveGroupInterface::Plan _plan)
{
        group.execute(_plan);
}

void arm_allfunction::planAndExecute()
{
        moveit::planning_interface::MoveItErrorCode success = group.plan(my_plan);
        if(success)
                group.execute(my_plan);
}

void arm_allfunction::yamlInit(const std::string _path)
{
        get_yaml_info.getYamlPath(_path);
        get_yaml_info.loadInfo();
}

void arm_allfunction::visualizeInit()
{
        visual_tools.deleteAllMarkers();
        visual_tools.loadRemoteControl();
}

void arm_allfunction::visualizePose(geometry_msgs::Pose  _pose)
{
        const robot_state::JointModelGroup* joint_model_group =
                group.getCurrentState()->getJointModelGroup(groupName);

        visual_tools.publishAxisLabeled(_pose, "pose");
        if(visual_tools.publishTrajectoryLine(my_plan.trajectory_,  joint_model_group))
                ROS_INFO("Vision ok");
        visual_tools.trigger();
}

void arm_allfunction::init(const std::string _path)
{
        visualizeInit();
        yamlInit(_path);
}

void arm_allfunction::updatePose(ChgDirectionEnumdef chg_direction, PoseDimensionEnumdef pose_dismension)
{
        switch (pose_dismension)
        {
                case PoseDimensionEnumdef::x:
                        if(chg_direction == ChgDirectionEnumdef::add)
                                pose.position.x += 0.01;
                        else
                                pose.position.x -= 0.01;
                        break;
                case PoseDimensionEnumdef::y:
                        if(chg_direction == ChgDirectionEnumdef::add)
                                pose.position.y += 0.01;
                        else
                                pose.position.y -= 0.01;
                        break;
                case PoseDimensionEnumdef::z:
                        if(chg_direction == ChgDirectionEnumdef::add)
                                pose.position.z += 0.01;
                        else
                                pose.position.z -= 0.01;
                        break;
                case PoseDimensionEnumdef::roll:
                        chgPoseOrientation(chg_direction, pose_dismension);
                        break;
                case PoseDimensionEnumdef::pitch:
                        chgPoseOrientation(chg_direction, pose_dismension);
                        break;
                case PoseDimensionEnumdef::yaw:
                        chgPoseOrientation(chg_direction, pose_dismension);
                        break;
                default:break;
        }
}

void arm_allfunction::chgPoseOrientation(ChgDirectionEnumdef chg_direction, PoseDimensionEnumdef pose_dismension)
{
        double r, p, y;
        tf::Quaternion qtn;
        tf::quaternionMsgToTF(pose.orientation, qtn);
        tf::Matrix3x3(qtn).getRPY(r, p, y);
        switch (pose_dismension)
        {
                case PoseDimensionEnumdef::roll:
                        if (chg_direction == ChgDirectionEnumdef::add)
                                r += 1.0 * 3.14 / 180;
                        else
                                r -= 1.0 * 3.14 / 180;
                        break;
                case PoseDimensionEnumdef::pitch:
                        if (chg_direction == ChgDirectionEnumdef::add)
                                p += 1.0 * 3.14 / 180;
                        else
                                p -= 1.0 * 3.14 / 180;
                        break;
                case PoseDimensionEnumdef::yaw:
                        if (chg_direction == ChgDirectionEnumdef::add)
                                y += 1.0 * 3.14 / 180;
                        else
                                y -= 1.0 * 3.14 / 180;
                        break;
                default:break;
        }
        pose.orientation = tf::createQuaternionMsgFromRollPitchYaw(r,p,y);
}

void arm_allfunction::printPoseInfo()
{       
        double r, p, y;
        tf::Quaternion qtn;
        tf::quaternionMsgToTF(pose.orientation, qtn);
        tf::Matrix3x3(qtn).getRPY(r, p, y);
        ROS_INFO("r: %.2f\r\n p: %.2f\r\n y:%.2f\r\n", r, p, y);
        ROS_INFO("x: %.2f\r\n y: %.2f\r\n z:%.2f\r\n", pose.position.x, pose.position.y, pose.position.z);        
}

moveit_msgs::RobotTrajectory arm_allfunction::turnPlanTime(moveit_msgs::RobotTrajectory trajectory_, double& scale)
{
        for (int j=0;j<trajectory_.joint_trajectory.points.size();j++) 
                for (int i=0;i<trajectory_.joint_trajectory.points[j].positions.size();i++) 
                        trajectory_.joint_trajectory.points[j].time_from_start.operator*=(scale);
        return trajectory_;
}

moveit_msgs::RobotTrajectory arm_allfunction::combineTrajectory(
        geometry_msgs::Pose msg1, 
        geometry_msgs::Pose msg2,
        moveit::planning_interface::MoveGroupInterface *move_group
)
{
        std::vector<moveit::planning_interface::MoveGroupInterface::Plan> my_double_plan;
        moveit::planning_interface::MoveItErrorCode success ;
        moveit_msgs::RobotTrajectory _trajectory;
        moveit_msgs::RobotTrajectory real_trajectory;

        moveit::planning_interface::MoveGroupInterface::Plan single_plan;
        moveit::core::RobotStatePtr start_state(group.getCurrentState());
        group.setStartStateToCurrentState();
        setTarget(msg1);               
       
        if(move_group->plan(single_plan)== moveit::planning_interface::MoveItErrorCode::SUCCESS)
                my_double_plan.push_back(single_plan);
        single_plan.start_state_.joint_state.position = single_plan.trajectory_.joint_trajectory.points.back().positions;
        move_group->setStartState(single_plan.start_state_);

        setTarget(msg2);
        success = move_group->plan(single_plan);
        if(move_group->plan(single_plan)== moveit::planning_interface::MoveItErrorCode::SUCCESS)
                my_double_plan.push_back(single_plan);

        single_plan.start_state_.joint_state.position = single_plan.trajectory_.joint_trajectory.points.back().positions;
        move_group->setStartState(single_plan.start_state_);
        

        _trajectory.joint_trajectory.joint_names = my_double_plan[0].trajectory_.joint_trajectory.joint_names;
        //_trajectory.joint_trajectory.points = my_double_plan[0].trajectory_.joint_trajectory.points;
        for (size_t i = 0; i < my_double_plan.size(); i++)
        {
                for (size_t j = 0; j < my_double_plan[i].trajectory_.joint_trajectory.points.size(); j++)
                {       
                        _trajectory.joint_trajectory.points.push_back(my_double_plan[i].trajectory_.joint_trajectory.points[j]);       
                }
        }

        real_trajectory = _trajectory;
        robot_trajectory::RobotTrajectory rt(move_group->getCurrentState()->getRobotModel(), move_group->getName());
        rt.setRobotTrajectoryMsg(*move_group->getCurrentState(), real_trajectory);
        trajectory_processing::IterativeParabolicTimeParameterization itpt;

        bool ItSuccess = itpt.computeTimeStamps(rt,1.0,1.0);
        ROS_INFO("Computed time stamp %s", ItSuccess ? "SUCCEDED" : "FAILED");

        rt.getRobotTrajectoryMsg(real_trajectory);

        return real_trajectory;
}
